#! /usr/bin/env python
# coding=utf-8

# 给定那个uniprot_tabel.tsv 表格 

from Bio import SeqIO
import re
import argparse
import sys
from argparse import RawTextHelpFormatter


parser = argparse.ArgumentParser(
    description='''
    对diamond 匹配了 swissprot的 蛋白 的m8文件 匹配到最高匹配的 swiss ID号 基因symbol 和 基因description
    用法:
    match_swiss_prot_annotation.py -i matches.tsv -s /home/good/lib/swissprot/uniprot_tabel.tsv -o out_gene2swiss_table
    ''',formatter_class=RawTextHelpFormatter)



parser.add_argument('-i',
                help='必须给定，输入的diamond m8文件 ')

parser.add_argument('-s',
                help='必须给定，uniprot 基因 id symbol description的文件')


parser.add_argument('-o',
                help='输出文件 一个表格')


# 取出pasa中被 merge的 mRNA

args = parser.parse_args()

if not args.i or not args.o or not args.s:
    parser.print_help()
    sys.exit()

swiss_prot_tab = args.s

out_gene2swiss_table = args.o

match_file = args.i


# swiss_prot_tab = '/home/good/lib/swissprot/uniprot_tabel.tsv'

# match_file = 'matches.tsv'

# out_gene2swiss_table = 'out_gene2swiss_table'

outfile = open(out_gene2swiss_table,'w')

# 读取 swiss tab
print('read swiss prot tab')
swiss_sym_dic = {}
swiss_des_dic = {}

with open(swiss_prot_tab) as fila:
	for i in fila:
		k = i.strip().split('\t')
		if len(k)>2:
			swiss_sym_dic[k[0]] = k[1]
			swiss_des_dic[k[0]] = k[2]

print('finish read')

gene2sym = {}
gene2des = {}

with open(match_file) as fila:
	for i in fila:
		k = i.strip().split('\t')
		if len(k)>1:
			k[1] = k[1].split('.')[0]

			if k[1] in swiss_sym_dic and k[0] not in gene2sym:
				gene2sym[k[0]] = swiss_sym_dic[k[1]]
				gene2des[k[0]] = swiss_des_dic[k[1]]
				outfile.write('\t'.join([k[0].split('.mRNA')[0],k[1],swiss_sym_dic[k[1]], swiss_des_dic[k[1]]])+'\n')
print('finish make table')

outfile.close()
